"""
    types.jl

Core type definitions for GSI analysis system, providing the fundamental
data structures used throughout the package.
"""

using LinearAlgebra
using StaticArrays

"""
    AbstractAnalysisConfig

Abstract base type for analysis configuration objects.
"""
abstract type AbstractAnalysisConfig end

# Normalize auxiliary configuration dictionaries to `Dict{String,Any}` so downstream
# code can consistently request string-valued keys regardless of the original input
# type (Symbols, Strings, NamedTuples, etc.).
_normalize_params(params::Nothing) = Dict{String,Any}()
function _normalize_params(params::AbstractDict)
    return Dict{String,Any}(string(k) => v for (k, v) in params)
end
function _normalize_params(params::NamedTuple)
    return Dict{String,Any}(string(k) => v for (k, v) in pairs(params))
end

"""
    AnalysisConfig{T<:AbstractFloat}

Configuration parameters for the GSI analysis system.

# Fields
- `grid_size::Tuple{Int,Int,Int}`: Grid dimensions (nx, ny, nz)
- `ensemble_size::Int`: Number of ensemble members for ensemble / hybrid methods
- `analysis_method::String`: High-level analysis driver tag (e.g. "3DVar", "4DVar", "EnKF", "Hybrid")
- `hybrid_coeff::T`: Hybrid ensemble coefficient (0.0 = static B, 1.0 = ensemble B)
- `precision::Type{T}`: Floating point precision (e.g. `Float32`, `Float64`)
- `use_hybrid::Bool`: Enable hybrid ensemble-variational method
- `max_iterations::Int`: Maximum inner-loop iterations
- `convergence_tol::T`: Convergence tolerance for minimization
- `model_type::String`: Underlying model identifier (e.g. "GFS", "WRF")
- `model_params::Dict{String,Any}`: Model-specific configuration options
- `output_path::String`: Destination for diagnostics and artifacts
- `params::Dict{String,Any}`: Additional workflow parameters (optimizer settings, run flags, etc.)
"""
struct AnalysisConfig{T<:AbstractFloat} <: AbstractAnalysisConfig
    grid_size::Tuple{Int,Int,Int}
    ensemble_size::Int
    analysis_method::String
    hybrid_coeff::T
    precision::Type{T}
    use_hybrid::Bool
    max_iterations::Int
    convergence_tol::T
    model_type::String
    model_params::Dict{String,Any}
    output_path::String
    params::Dict{String,Any}
    
    function AnalysisConfig{T}(;
        grid_size::Tuple{Int,Int,Int} = (360, 180, 64),
        ensemble_size::Int = 40,
        analysis_method::AbstractString = "3DVar",
        hybrid_coeff::Real = 0.75,
        precision::Type{T} = T,
        use_hybrid::Bool = true,
        max_iterations::Int = 200,
        convergence_tol::Real = 1e-5,
        model_type::AbstractString = "GFS",
        model_params::Union{Nothing,AbstractDict,NamedTuple} = nothing,
        output_path::AbstractString = "analysis_output",
        params::Union{Nothing,AbstractDict,NamedTuple} = nothing,
        kwargs...
    ) where T<:AbstractFloat
        hybrid_coeff_T = convert(T, hybrid_coeff)
        convergence_tol_T = convert(T, convergence_tol)
        model_params_dict = _normalize_params(model_params)
        params_dict = _normalize_params(params)
        if !isempty(kwargs)
            merge!(params_dict, Dict{String,Any}(string(k) => v for (k, v) in kwargs))
        end
        new{T}(
            grid_size,
            ensemble_size,
            String(analysis_method),
            hybrid_coeff_T,
            precision,
            use_hybrid,
            max_iterations,
            convergence_tol_T,
            String(model_type),
            model_params_dict,
            String(output_path),
            params_dict,
        )
    end
end

# Convenience constructor
AnalysisConfig(args...; kwargs...) = AnalysisConfig{Float64}(args...; kwargs...)

"""
    AbstractControlVector{T}

Abstract base type for control vectors in GSI analysis.
"""
abstract type AbstractControlVector{T<:AbstractFloat} end

"""
    AbstractStateVector{T}

Abstract base type for state vectors in GSI analysis.
"""
abstract type AbstractStateVector{T<:AbstractFloat} end

"""
    AbstractGrid

Abstract base type for grid structures in GSI analysis.
"""
abstract type AbstractGrid end

"""
    GridDimensions

Structure to hold grid dimension information.
"""
struct GridDimensions
    nx::Int  # Longitude grid points
    ny::Int  # Latitude grid points  
    nz::Int  # Vertical levels
    total::Int
    
    GridDimensions(nx::Int, ny::Int, nz::Int) = new(nx, ny, nz, nx * ny * nz)
end

GridDimensions(grid_size::Tuple{Int,Int,Int}) = GridDimensions(grid_size...)

"""
    VariableInfo

Information about a single analysis variable.
"""
struct VariableInfo
    name::String
    levels::Int
    start_index::Int
    end_index::Int
    has_bias::Bool
end

"""
    AnalysisVariables

Container for all analysis variables and their metadata.
"""
struct AnalysisVariables
    variables::Vector{VariableInfo}
    total_size::Int
    
    function AnalysisVariables(var_configs::Vector{Tuple{String,Int,Bool}})
        variables = VariableInfo[]
        current_index = 1
        
        for (name, levels, has_bias) in var_configs
            var_size = levels
            var_info = VariableInfo(name, levels, current_index, current_index + var_size - 1, has_bias)
            push!(variables, var_info)
            current_index += var_size
        end
        
        new(variables, current_index - 1)
    end
end

"""
    GridConfiguration
    
A temporary alias for AbstractAnalysisConfig to maintain compatibility with EnKF modules.
This will be properly defined when GridOperations is loaded.
"""
const GridConfiguration = AbstractAnalysisConfig

"""
    default_analysis_variables()

Returns default GSI analysis variable configuration.
"""
function default_analysis_variables()
    var_configs = [
        ("u", 64, false),      # Zonal wind
        ("v", 64, false),      # Meridional wind  
        ("t", 64, false),      # Temperature
        ("q", 64, false),      # Specific humidity
        ("ps", 1, false),      # Surface pressure
        ("oz", 64, false),     # Ozone
        ("cw", 64, false),     # Cloud water
        ("stl", 1, false),     # Skin temperature over land
        ("sti", 1, false)      # Skin temperature over ice
    ]
    return AnalysisVariables(var_configs)
end
